#!/usr/bin/python3
import os
import shutil
from pathlib import Path

def clean_knapsack_data(input_file, output_file):
    """
    清理背包数据文件，移除大规模数据集中的解向量行
    """
    with open(input_file, 'r') as f:
        lines = f.readlines()
    
    # 读取第一行获取物品数量
    first_line = lines[0].strip().split()
    n = int(first_line[0])  # 物品数量
    
    # 标准UU格式应该有 n+1 行（第一行是n和capacity，后面n行是物品）
    expected_lines = n + 1
    
    # 如果行数超过预期，说明有额外的解向量，需要移除
    if len(lines) > expected_lines:
        print(f"    检测到额外行数，移除解向量: {len(lines)} -> {expected_lines}")
        lines = lines[:expected_lines]
    
    # 写入清理后的数据
    with open(output_file, 'w') as f:
        f.writelines(lines)

def convert_dataset_structure(source_dir, target_dir):
    """
    将UU数据集格式转换为run_tests.py期望的格式，并清理数据
    """
    source_path = Path(source_dir)
    target_path = Path(target_dir)
    
    # 确保目标目录存在
    target_path.mkdir(parents=True, exist_ok=True)
    
    # 处理的子目录
    subdirs = ['low-dimensional', 'large_scale']
    
    for subdir in subdirs:
        subdir_path = source_path / subdir
        if not subdir_path.exists():
            print(f"子目录不存在: {subdir_path}")
            continue
        
        print(f"处理子目录: {subdir}")
        
        # 遍历子目录中的所有文件
        for file_path in subdir_path.iterdir():
            if file_path.is_file():
                # 创建以文件名命名的目录
                file_dir = target_path / file_path.name
                file_dir.mkdir(exist_ok=True)
                
                # 目标文件路径
                target_file = file_dir / "test.in"
                
                # 根据文件类型决定是否需要清理
                if subdir == 'large_scale':
                    # 大规模数据集需要清理解向量
                    print(f"  转换并清理: {file_path} -> {target_file}")
                    clean_knapsack_data(file_path, target_file)
                else:
                    # 小规模数据集直接复制
                    print(f"  转换: {file_path} -> {target_file}")
                    shutil.copy2(file_path, target_file)

if __name__ == "__main__":
    source_dir = "./data"  # 当前数据集目录（修复路径）
    target_dir = "./data_converted_fixed"  # 修复后的目录
    
    print("开始转换数据集目录结构（修复版）...")
    convert_dataset_structure(source_dir, target_dir)
    print(f"转换完成！修复后的数据集保存在: {target_dir}")
    print("大规模数据集的解向量已被移除") 